# 计算模型的时间复杂度
# 参考：https://blog.csdn.net/weixin_41519463/article/details/102468868
import PIL.Image as pil_image
# lr_image_path = r'C:\lcj\02-dataset\920_experi_beads\0919_thy1\test\52\52_lr.bmp'
# lr = pil_image.open(lr_image_path)
from ptflops import get_model_complexity_info
from models import ResAttU_Net_noBN, UNet, VDSR, BNAnet

net = UNet()
flops, params = get_model_complexity_info(net, (1, 1000, 1000), as_strings=True, print_per_layer_stat=True)
print('Flops:  ' + flops)
print('Params: ' + params)